#include "eth.h"
#include "ip.h"
#include "arp.h"
#include "util.h"
#include "udp.h"
#include "posix_api.h"


void _encode_udp_hdr_pkt(uint8_t * msg, uint16_t src_port, uint16_t dst_port, uint16_t total_len, unsigned char * data){
    struct rte_ipv4_hdr * iphdr = (struct rte_ipv4_hdr *)(msg + sizeof(struct rte_ether_hdr));
    struct rte_udp_hdr * udphdr = (struct rte_udp_hdr *)(iphdr + 1);
    udphdr->src_port = src_port;
    udphdr->dst_port = dst_port;
    uint16_t udp_len = total_len - sizeof(struct rte_ether_hdr) - sizeof(struct rte_ipv4_hdr);
    udphdr->dgram_len = htons(udp_len);
    rte_memcpy((uint8_t*)(udphdr + 1), data, udp_len);
    udphdr->dgram_cksum = 0;
    udphdr->dgram_cksum = rte_ipv4_udptcp_cksum(iphdr, udphdr);
    print_address("udp  --> src :  ", iphdr->src_addr);
    printf("src port : %d \n", ntohs(src_port));
    print_address("udp  --> dst :  ", iphdr->dst_addr);
    printf("dst port : %d \n", ntohs(dst_port));
}


struct rte_mbuf * _encode_udp(uint8_t * msg, uint32_t src_ip, uint32_t dst_ip, uint16_t src_port, uint16_t dst_port,
    uint8_t * src_mac, uint8_t * dst_mac, uint8_t * data, uint16_t length){

    encode_ether_hdr_pkt(msg, dst_mac, src_mac);
    encode_ip_hdr_pkt(msg, length - sizeof (struct rte_ether_hdr), IPPROTO_UDP, src_ip, dst_ip);
    _encode_udp_hdr_pkt(msg, src_port, dst_port, length, data);
}

struct rte_mbuf * _send_udp(struct rte_mempool * pool, uint32_t src_ip, uint32_t dst_ip, uint16_t src_port, uint16_t dst_port,
        uint8_t * src_mac, uint8_t * dst_mac, uint8_t * data, uint16_t length){
    const unsigned total_len = length + 42;
    struct rte_mbuf * m_buf = rte_pktmbuf_alloc(pool);
    if(!m_buf){
        error_print("__send_udp");
    }
    m_buf->data_len = total_len;
    m_buf->pkt_len = total_len;
    uint8_t * pkt_data = rte_pktmbuf_mtod(m_buf, uint8_t *);
    _encode_udp(pkt_data, src_ip, dst_ip, src_port, dst_port, src_mac, dst_mac, data, length);
    return m_buf;
}

int udp_pkt_process(struct rte_mbuf * m_buf, struct rte_mempool * pool){
    struct inout_ring * ring = ring_instance();
    struct rte_ether_hdr * ethdr = rte_pktmbuf_mtod(m_buf,struct rte_ether_hdr *);
    struct rte_ipv4_hdr *iphdr =  (struct rte_ipv4_hdr *)(ethdr + 1);
    struct rte_udp_hdr *udphdr = (struct rte_udp_hdr *)(iphdr + 1);
    rte_memcpy(g_dst_mac, ethdr->src_addr.addr_bytes, RTE_ETHER_ADDR_LEN);

    rte_memcpy(&g_src_ip, &iphdr->dst_addr, sizeof(uint32_t));
    rte_memcpy(&g_dst_ip, &iphdr->src_addr, sizeof(uint32_t));

    rte_memcpy(&g_src_port, &udphdr->dst_port, sizeof(uint16_t));
    rte_memcpy(&g_dst_port, &udphdr->src_port, sizeof(uint16_t));

    uint16_t length = ntohs(udphdr->dgram_len);
    *((char*)udphdr + length) = '\0';

    print_address("__udp_process src : ", iphdr->src_addr);
    printf(" __udp_process src prot: %d \n", ntohs(udphdr->src_port));
    print_address("__udp_process dsc : ", iphdr->dst_addr);
    printf("__udp_process dst prot: %d  context: %s \n", ntohs(udphdr->dst_port), (char *)(udphdr+1));

    struct rte_mbuf * txbuf = _send_udp(pool, g_src_ip, g_dst_ip, g_src_port, g_dst_port, g_src_mac, g_dst_mac, (uint8_t *)(udphdr+1), length);
    rte_ring_mp_enqueue_burst(ring->out, (void **)&txbuf, 1, NULL);
    rte_pktmbuf_free(txbuf);
}


int udp_process(struct rte_mbuf * m_buf){
    struct rte_ipv4_hdr * iphdr = rte_pktmbuf_mtod_offset(m_buf, struct rte_ipv4_hdr *, sizeof(struct rte_ether_hdr));
    struct rte_udp_hdr * udphdr = (struct rte_udp_hdr *)(iphdr + 1);
    print_address("udp process ----> src :  ", iphdr->src_addr);
    printf("%d \n", udphdr->src_port);

    struct local_host * host = get_host_info_from_ip_port(iphdr->dst_addr, udphdr->dst_port, iphdr->next_proto_id);
    if(host == NULL){
        rte_pktmbuf_free(m_buf);
        return -3;
    }

    struct udp_pkg * pkg = rte_malloc("udp pkg", sizeof(struct udp_pkg), 0);
    if(pkg == NULL){
        rte_pktmbuf_free(m_buf);
        return  -1;
    }
    memset(pkg, 0, sizeof(struct udp_pkg));

    pkg->d_ip = iphdr->dst_addr;
    pkg->s_ip = iphdr->src_addr;
    pkg->s_port = udphdr->src_port;
    pkg->d_port = udphdr->dst_port;

    pkg->protocol = IPPROTO_UDP;
    pkg->length = ntohs(udphdr->dgram_len);

    pkg->data = rte_malloc("udp data ", pkg->length - sizeof(struct rte_udp_hdr), 0);
    if(pkg->data == NULL){
        rte_pktmbuf_free(m_buf);
        rte_free(pkg);
        return  -2;
    }
    rte_memcpy(pkg->data, (unsigned char *)(udphdr + 1), pkg->length - sizeof(struct rte_udp_hdr));
    rte_ring_mp_enqueue(host->recv_buf, pkg);
    pthread_mutex_lock(&host->mtx);
    pthread_cond_signal(&host->cond);
    pthread_mutex_unlock(&host->mtx);
    rte_pktmbuf_free(m_buf);
    return 0;
}


int udp_server_entry(__attribute__((unused)) void * arg){
    int conn_fd = m_socket(AF_INET, SOCK_DGRAM, 0);
    if(conn_fd == -1){
        printf("socket failed \n");
        return  -1;
    }

    struct sockaddr_in server_address, client_address;
    memset(&server_address, 0, sizeof(struct sockaddr_in));
    server_address.sin_family = AF_INET;
    server_address.sin_port = htons(G_PORT);
    server_address.sin_addr.s_addr = g_local_ip;

    m_bind(conn_fd, (struct sockaddr *)&server_address, sizeof(struct sockaddr_in));
    char buffer[UDP_APP_RECV_BUFFER_SIZE] = {0};
    socklen_t  len = sizeof(client_address);
    while(1){
        if(m_recv_from(conn_fd, buffer, UDP_APP_RECV_BUFFER_SIZE, 0, (struct sockaddr *)&client_address, &len) < 0){
            continue;
        } else{
            print_address("recv from : ", client_address.sin_addr.s_addr);
            printf("prot : %d ; data:  %s \n", client_address.sin_port, buffer);
            m_send_to(conn_fd, buffer, strlen(buffer), 0, (struct sockaddr *)&client_address, sizeof(struct sockaddr ));
        }
    }
    m_close(conn_fd);
}


int udp_out(struct rte_mempool * pool){
    struct local_host * host;
    for(host = g_hosts; host != NULL; host = host->next){
        struct udp_pkg * pkg;
        int l_send = rte_ring_mc_dequeue(host->send_buf, (void **)&pkg);
        if(l_send < 0){
            continue;
        }

        print_address("udp out src :", pkg->d_ip);
        printf("port %d \n", pkg->d_port);
        printf("udp out after arp: %s \n ", pkg->data);

        uint8_t  * dst_mac = get_dst_mac_address(pkg->d_ip);
        if(dst_mac == NULL){
            struct rte_mbuf * arp_buf = arp_send(pool, RTE_ARP_OP_REQUEST, g_default_arp_mac, pkg->s_ip, pkg->d_ip);
            struct inout_ring * ring = ring_instance();
            rte_ring_mp_enqueue_burst(ring->out, (void **)&arp_buf, 1, NULL);
            rte_ring_mp_enqueue(host->send_buf, pkg);
        }else{
            struct rte_mbuf * udp_buf = _send_udp(pool, pkg->s_ip, pkg->d_ip, pkg->s_port, pkg->d_port, host->local_mac, dst_mac, pkg->data, pkg->length);
            struct inout_ring * ring = ring_instance();
            rte_ring_mp_enqueue_burst(ring->out, (void **)&udp_buf, 1, NULL);
        }
    }
    return  0;
}